import os
import sys
import numpy as np

class Topo:
    def _get_layer_rank(self, rank_size, rank_stride, level_rank):
        if rank_stride > level_rank:
            return 1
        else:
            # return rank_size if level_rank > rank_size*rank_stride else level_rank // rank_stride
            if level_rank > rank_size*rank_stride:
                return rank_size
            else:
                return level_rank // rank_stride
        
    def _get_l0layer_rank(self, rank_size, rank_stride):
        return self._get_layer_rank(rank_size, rank_stride, 2)

    def _get_l1layer_rank(self, rank_size, rank_stride):
        l0 = self._get_layer_rank(rank_size, rank_stride, 2)
        l1 = self._get_layer_rank(rank_size, rank_stride, 2*8)
        return l1 // l0

    def _get_l2layer_rank(self, rank_size, rank_stride):
        l1 = self._get_layer_rank(rank_size, rank_stride, 2*8)
        l2 = self._get_layer_rank(rank_size, rank_stride, 2*8*48)
        return l2 // l1

    def _get_l3layer_rank(self, rank_size, rank_stride):
        l2 = self._get_layer_rank(rank_size, rank_stride, 2*8*48)
        return rank_size // l2

    def domain_to_topolayer(self, rank_size, rank_stride):
        topolayer = list()
        topolayer.append(self._get_l0layer_rank(rank_size, rank_stride))
        topolayer.append(self._get_l1layer_rank(rank_size, rank_stride))
        topolayer.append(self._get_l2layer_rank(rank_size, rank_stride))
        topolayer.append(self._get_l3layer_rank(rank_size, rank_stride))

        print(topolayer)
        return topolayer

topo = Topo()
topo.domain_to_topolayer(8, 1)
topo.domain_to_topolayer(64, 1)
topo.domain_to_topolayer(128, 2)
topo.domain_to_topolayer(256, 2)
topo.domain_to_topolayer(768, 2)

print(f"====")
topo.domain_to_topolayer(16, 96)
topo.domain_to_topolayer(2, 768)
topo.domain_to_topolayer(16, 8)
topo.domain_to_topolayer(48, 16)
topo.domain_to_topolayer(96, 16)

topo.domain_to_topolayer(8, 2)
topo.domain_to_topolayer(16, 4)
topo.domain_to_topolayer(16, 8)
topo.domain_to_topolayer(16, 16)